###############################################################################   
#### Replication Materials                                                 #### 
#### Kim, Nakka, Gopal, Desmrais, Mancinelli, Harden, Ko, Boehmke. 2021.   ####
#### Attention to the COVID-19 pandemic on Twitter:                        ####
#### Partisan differences among U.S. state legislators                     ####
#### Legislative Studies Quarterly                                         ####
###############################################################################  


###############################################################################
################################### Set Up ####################################
###############################################################################

# packages -------------------------

lapply(c('readr', 'stargazer', 'texreg', 'xtable', 'lme4', 'glmmTMB','plm',
  'lmtest', 'interplot','dummies','effects', 'ggthemes','caret', 
  'ggeffects', 'scales','ggrepel','sandwich','reshape2'), 
  require, 
  character.only = TRUE
  )

# read data set for regression -------------------------

data_path <- '/Users/taegyoon/Google Drive/spap_state/spap_state_attention/data/' 
df <- read_csv(paste0(data_path, "spap_state_attention_regression.csv")) 
df <- subset(df, select = -c(X1) )


###############################################################################
############################ Fit Regression Model #############################
###############################################################################

# set the panel structure -------------------------

plm_df <- pdata.frame(
  x = df, 
  index = c('user.screen_name', 'week'))

# fit models -------------------------

main_all_inter <- plm(covid_relevant_1_log ~ 
                        party_new*state_case_pop + 
                        party_new*state_death_pop +
                        party_new*national_case_pop + 
                        party_new*national_death_pop +
                        party_new*state_covid_policy + 
                        state_abbrev +
                        week_new + week_new_2 + week_new_3,            
                      data = plm_df,
                      index = c('user.screen_name', 'week_new'),
                      model = 'random', 
                      effect = 'twoways')


###############################################################################
####################### Generate Effect Plots: Figure 2 #######################
###############################################################################

# variance-covariance matrix for adjusted SE -------------------------

vcov <- vcovHC(main_all_inter, 
               method = 'arellano')

# state case -------------------------

state_case_values <- seq(round(min(df$state_case_pop), 1),
                         round(max(df$state_case_pop), 1),
                         0.1)
state_case_pop_effect <- Effect(focal.predictors = c('state_case_pop',
                                                     'party_new'), 
                                mod = main_all_inter,
                                xlevels = list(state_case_pop = state_case_values),
                                given.values=c(
                                  state_abbrevCA = 0,
                                  state_abbrevTX = 0,
                                  state_abbrevAZ = 0,
                                  state_abbrevNY = 0,
                                  state_abbrevPA = 0,
                                  state_abbrevMA = 0,
                                  state_abbrevTN = 0,
                                  state_abbrevNM = 0,
                                  state_abbrevFL = 0,
                                  state_abbrevVA = 0,
                                  state_abbrevNC = 0,
                                  state_abbrevMN = 0,
                                  state_abbrevMD = 0,
                                  state_abbrevUT = 0,
                                  state_abbrevMI = 0,
                                  state_abbrevOH = 0,
                                  state_abbrevIA = 0,
                                  state_abbrevMO = 0,
                                  state_abbrevNJ = 0,
                                  state_abbrevNH = 0,
                                  state_abbrevWI = 0,
                                  state_abbrevAR = 0,
                                  state_abbrevOR = 0,
                                  state_abbrevGA = 0,
                                  state_abbrevCO = 1,
                                  state_abbrevSC = 0,
                                  state_abbrevCT = 0,
                                  state_abbrevOK = 0,
                                  state_abbrevAL = 0,
                                  state_abbrevDE = 0,
                                  state_abbrevNV = 0,
                                  state_abbrevKY = 0,
                                  state_abbrevIN = 0,
                                  state_abbrevND = 0,
                                  state_abbrevIL = 0,
                                  state_abbrevKS = 0,
                                  state_abbrevWA = 0,
                                  state_abbrevMS = 0,
                                  state_abbrevWV = 0,
                                  state_abbrevLA = 0,
                                  state_abbrevRI = 0,
                                  state_abbrevSD = 0,
                                  state_abbrevWY = 0,
                                  state_abbrevID = 0,
                                  state_abbrevHI = 0,
                                  state_abbrevMT = 0,
                                  state_abbrevME = 0,
                                  state_abbrevVT = 0,
                                  state_death_pop = median(df$state_death_pop),
                                  national_case_pop = median(df$national_case_pop),
                                  national_death_pop = median(df$national_death_pop),
                                  state_covid_policy = median(df$state_covid_policy),
                                  week_new = median(df$week_new),
                                  week_new_2 = median(df$week_new_2),
                                  week_new_3 = median(df$week_new_3)))

state_case_me_se_dem <- sqrt(vcov[4, 4]) # SE for marginal effect of predictor
state_case_me_se_other <- sqrt(vcov[4, 4] 
                               + (1)^2 * vcov[60, 60] + 
                                 2 * 1 * vcov[4, 60]) # SE for marginal effect of predictor
state_case_me_se_rep <- sqrt(vcov[4, 4] 
                             + (1)^2 * vcov[61, 61] + 
                               2 * 1 * vcov[4,61]) # SE for marginal effect of predictor

state_case_me_ci_dem <- cbind(state_case_values * (1.96*state_case_me_se_dem), 
                              state_case_values * -(1.96*state_case_me_se_dem))

state_case_me_ci_other <- cbind(state_case_values * (1.96*state_case_me_se_other), 
                                state_case_values * -(1.96*state_case_me_se_other))

state_case_me_ci_rep <- cbind(state_case_values * (1.96*state_case_me_se_rep), 
                              state_case_values * -(1.96*state_case_me_se_rep))

state_case_preds <- state_case_pop_effect$fit # in the order of Dem, Neither, Rep
state_case_upper <- state_case_pop_effect$fit + 
  c(state_case_me_ci_dem[,1], state_case_me_ci_other[,1], state_case_me_ci_rep[,1])
state_case_lower <- state_case_pop_effect$fit + 
  c(state_case_me_ci_dem[,2], state_case_me_ci_other[,2], state_case_me_ci_rep[,2])
state_case_final <- data.frame(Party = c(rep('Democrat',
                                             length(state_case_values)), 
                                         rep('Other',
                                             length(state_case_values)),
                                         rep('Republican',
                                             length(state_case_values))), 
                               X = as.numeric(state_case_values), 
                               Y = as.numeric(state_case_preds), 
                               U = as.numeric(state_case_upper), 
                               L = as.numeric(state_case_lower),
                               Y_linear = exp(as.numeric(state_case_preds)) - 1, 
                               U_linear = exp(as.numeric(state_case_upper)) - 1, 
                               L_linear = exp(as.numeric(state_case_lower)) - 1)

state_case_final$Party <- as.factor(state_case_final$Party)
state_case_final$Party <- factor(state_case_final$Party, 
                                 levels = rev(levels(state_case_final$Party)))
state_case_final <- state_case_final[which(state_case_final$Party != 'Other'), ]

ggplot(state_case_final,
       aes(x = X, 
           y = Y_linear, 
           color = Party)) + 
  theme_calc() +
  theme(text = element_text(size = 12.5), 
        plot.title = element_text(hjust = 0.5),
        plot.margin = grid::unit(c(5, 5, 5, 5), "mm"),
        legend.title = element_blank(),
        legend.position='bottom',
        legend.direction='horizontal') +
  geom_line() +
  scale_color_manual(values = c('darkred', 'dodgerblue4')) +
  geom_ribbon(aes(ymin = L_linear, 
                  ymax = U_linear), 
              linetype = 2, 
              alpha = 0) +
  xlab('\nWeekly Count of State New Cases (per 10k)') + 
  ylab('Count of COVID-19 Tweets\n') +
  geom_rug(data = df, 
           aes(x = state_case_pop), 
           inherit.aes = FALSE, 
           color = 'gray60') +
  scale_shape_cleveland(overlap = FALSE)

ggsave(paste0(rep_path, 'Fig2a.png'),
       dpi = 600,
       width = 5,
       height = 3.5,
       units = 'in')

# state death -------------------------

state_death_values <- seq(round(min(df$state_death_pop),1),
                          round(max(df$state_death_pop),1),
                          0.1)
state_death_pop_effect <- Effect(focal.predictors = c('state_death_pop',
                                                      'party_new'), 
                                 mod = main_all_inter,
                                 xlevels = list(state_death_pop = state_death_values),
                                 given.values=c(
                                   state_abbrevCA = 0,
                                   state_abbrevTX = 0,
                                   state_abbrevAZ = 0,
                                   state_abbrevNY = 0,
                                   state_abbrevPA = 0,
                                   state_abbrevMA = 0,
                                   state_abbrevTN = 0,
                                   state_abbrevNM = 0,
                                   state_abbrevFL = 0,
                                   state_abbrevVA = 0,
                                   state_abbrevNC = 0,
                                   state_abbrevMN = 0,
                                   state_abbrevMD = 0,
                                   state_abbrevUT = 0,
                                   state_abbrevMI = 0,
                                   state_abbrevOH = 0,
                                   state_abbrevIA = 0,
                                   state_abbrevMO = 0,
                                   state_abbrevNJ = 0,
                                   state_abbrevNH = 0,
                                   state_abbrevWI = 0,
                                   state_abbrevAR = 0,
                                   state_abbrevOR = 0,
                                   state_abbrevGA = 0,
                                   state_abbrevCO = 1,
                                   state_abbrevSC = 0,
                                   state_abbrevCT = 0,
                                   state_abbrevOK = 0,
                                   state_abbrevAL = 0,
                                   state_abbrevDE = 0,
                                   state_abbrevNV = 0,
                                   state_abbrevKY = 0,
                                   state_abbrevIN = 0,
                                   state_abbrevND = 0,
                                   state_abbrevIL = 0,
                                   state_abbrevKS = 0,
                                   state_abbrevWA = 0,
                                   state_abbrevMS = 0,
                                   state_abbrevWV = 0,
                                   state_abbrevLA = 0,
                                   state_abbrevRI = 0,
                                   state_abbrevSD = 0,
                                   state_abbrevWY = 0,
                                   state_abbrevID = 0,
                                   state_abbrevHI = 0,
                                   state_abbrevMT = 0,
                                   state_abbrevME = 0,
                                   state_abbrevVT = 0,
                                   state_case_pop = median(df$state_case_pop),
                                   national_case_pop = median(df$national_case_pop),
                                   national_death_pop = median(df$national_death_pop),
                                   state_covid_policy = median(df$state_covid_policy),
                                   week_new = median(df$week_new),
                                   week_new_2 = median(df$week_new_2),
                                   week_new_3 = median(df$week_new_3)))

state_death_me_se_dem <- sqrt(vcov[5, 5]) # SE for marginal effect of predictor
state_death_me_se_other <- sqrt(vcov[5, 5] + 
                                  (1) ^ 2 * vcov[62, 62] + 
                                  2 * 1 * vcov[5, 62]) # SE for marginal effect of predictor
state_death_me_se_rep <- sqrt(vcov[5,5] + 
                                (1) ^ 2 * vcov[63, 63] + 
                                2 * 1 * vcov[5, 63]) # SE for marginal effect of predictor

state_death_me_ci_dem <- cbind(state_death_values * (1.96*state_death_me_se_dem), 
                               state_death_values * -(1.96*state_death_me_se_dem))
state_death_me_ci_other <- cbind(state_death_values * (1.96*state_death_me_se_other), 
                                 state_death_values * -(1.96*state_death_me_se_other))
state_death_me_ci_rep <- cbind(state_death_values * (1.96*state_death_me_se_rep), 
                               state_death_values * -(1.96*state_death_me_se_rep))

state_death_preds <- state_death_pop_effect$fit # in the order of Dem, Neither, Rep
state_death_upper <- state_death_pop_effect$fit + c(state_death_me_ci_dem[,1], 
                                                    state_death_me_ci_other[,1], state_death_me_ci_rep[,1])
state_death_lower <- state_death_pop_effect$fit + c(state_death_me_ci_dem[,2], 
                                                    state_death_me_ci_other[,2], state_death_me_ci_rep[,2])
state_death_final <- data.frame(Party = c(rep('Democrat',
                                              length(state_death_values)), 
                                          rep('Other',
                                              length(state_death_values)),
                                          rep('Republican',
                                              length(state_death_values))), 
                                X = as.numeric(state_death_values), 
                                Y = as.numeric(state_death_preds), 
                                U = as.numeric(state_death_upper), 
                                L = as.numeric(state_death_lower),
                                Y_linear = exp(as.numeric(state_death_preds)) - 1, 
                                U_linear = exp(as.numeric(state_death_upper)) - 1, 
                                L_linear = exp(as.numeric(state_death_lower)) - 1)
state_death_final$Party <- as.factor(state_death_final$Party)
state_death_final$Party <- factor(state_death_final$Party, levels = rev(levels(state_death_final$Party)))
state_death_final <- state_death_final[which(state_death_final$Party!='Other'),]

ggplot(state_death_final, 
       aes(x = X, 
           y = Y_linear, 
           color = Party)) + 
  theme_calc() +
  theme(text = element_text(size = 12.5), 
        plot.title = element_text(hjust = 0.5),
        plot.margin = grid::unit(c(5, 5, 5, 5), "mm"),
        legend.title = element_blank(),
        legend.position='bottom',
        legend.direction='horizontal') +
  geom_line() +
  scale_color_manual(values = c('darkred', 'dodgerblue4')) +
  geom_ribbon(aes(ymin = L_linear, 
                  ymax = U_linear), 
              linetype = 2, 
              alpha = 0.0) +
  xlab('\nWeekly Count of State New Deaths (per 10k)') + 
  ylab('Count of COVID-19 Tweets\n') +
  geom_rug(data = df, 
           aes(x = state_death_pop), 
           inherit.aes = FALSE, 
           color = 'gray60') +
  scale_shape_cleveland(overlap = FALSE)

ggsave(paste0(rep_path, 'Fig2b.png'),
       dpi = 600,
       width = 5,
       height = 3.5,
       units = 'in')


# national case -------------------------

national_case_values <- seq(round(min(df$national_case_pop), 1),
                            round(max(df$national_case_pop), 1),
                            0.1)
national_case_pop_effect <- Effect(focal.predictors = c('national_case_pop','party_new'), 
                                   mod = main_all_inter,
                                   xlevels = list(national_case_pop = national_case_values),
                                   given.values=c(
                                     state_abbrevCA = 0,
                                     state_abbrevTX = 0,
                                     state_abbrevAZ = 0,
                                     state_abbrevNY = 0,
                                     state_abbrevPA = 0,
                                     state_abbrevMA = 0,
                                     state_abbrevTN = 0,
                                     state_abbrevNM = 0,
                                     state_abbrevFL = 0,
                                     state_abbrevVA = 0,
                                     state_abbrevNC = 0,
                                     state_abbrevMN = 0,
                                     state_abbrevMD = 0,
                                     state_abbrevUT = 0,
                                     state_abbrevMI = 0,
                                     state_abbrevOH = 0,
                                     state_abbrevIA = 0,
                                     state_abbrevMO = 0,
                                     state_abbrevNJ = 0,
                                     state_abbrevNH = 0,
                                     state_abbrevWI = 0,
                                     state_abbrevAR = 0,
                                     state_abbrevOR = 0,
                                     state_abbrevGA = 0,
                                     state_abbrevCO = 1,
                                     state_abbrevSC = 0,
                                     state_abbrevCT = 0,
                                     state_abbrevOK = 0,
                                     state_abbrevAL = 0,
                                     state_abbrevDE = 0,
                                     state_abbrevNV = 0,
                                     state_abbrevKY = 0,
                                     state_abbrevIN = 0,
                                     state_abbrevND = 0,
                                     state_abbrevIL = 0,
                                     state_abbrevKS = 0,
                                     state_abbrevWA = 0,
                                     state_abbrevMS = 0,
                                     state_abbrevWV = 0,
                                     state_abbrevLA = 0,
                                     state_abbrevRI = 0,
                                     state_abbrevSD = 0,
                                     state_abbrevWY = 0,
                                     state_abbrevID = 0,
                                     state_abbrevHI = 0,
                                     state_abbrevMT = 0,
                                     state_abbrevME = 0,
                                     state_abbrevVT = 0,
                                     state_case_pop = median(df$state_case_pop),
                                     state_death_pop = median(df$state_death_pop),
                                     national_death_pop = median(df$national_death_pop),
                                     state_covid_policy = median(df$state_covid_policy),
                                     week_new = median(df$week_new),
                                     week_new_2 = median(df$week_new_2),
                                     week_new_3 = median(df$week_new_3)))

national_case_me_se_dem <- sqrt(vcov[6, 6]) # SE for marginal effect of predictor
national_case_me_se_other <- sqrt(vcov[6, 6] 
                                  + (1) ^ 2 * vcov[64, 64] + 
                                    2 * 1 * vcov[6, 64]) # SE for marginal effect of predictor
national_case_me_se_rep <- sqrt(vcov[6, 6] 
                                + (1) ^ 2 * vcov[65, 65] + 
                                  2 * 1 * vcov[6, 65]) # SE for marginal effect of predictor

national_case_me_ci_dem <- cbind(national_case_values * (1.96*national_case_me_se_dem), 
                                 national_case_values * -(1.96*national_case_me_se_dem))
national_case_me_ci_other <- cbind(national_case_values * (1.96*national_case_me_se_other), 
                                   national_case_values * -(1.96*national_case_me_se_other))
national_case_me_ci_rep <- cbind(national_case_values * (1.96*national_case_me_se_rep), 
                                 national_case_values * -(1.96*national_case_me_se_rep))

national_case_preds <- national_case_pop_effect$fit # in the order of Dem, Neither, Rep
national_case_upper <- national_case_pop_effect$fit + c(national_case_me_ci_dem[,1], 
                                                        national_case_me_ci_other[,1], 
                                                        national_case_me_ci_rep[,1])
national_case_lower <- national_case_pop_effect$fit + c(national_case_me_ci_dem[,2], 
                                                        national_case_me_ci_other[,2], 
                                                        national_case_me_ci_rep[,2])
national_case_final <- data.frame(Party = c(rep('Democrat',
                                                length(national_case_values)), 
                                            rep('Other',
                                                length(national_case_values)),
                                            rep('Republican',
                                                length(national_case_values))), 
                                  X = as.numeric(national_case_values), 
                                  Y = as.numeric(national_case_preds), 
                                  U = as.numeric(national_case_upper), 
                                  L = as.numeric(national_case_lower),
                                  Y_linear = exp(as.numeric(national_case_preds)) - 1, 
                                  U_linear = exp(as.numeric(national_case_upper)) - 1, 
                                  L_linear = exp(as.numeric(national_case_lower)) - 1)

national_case_final$Party <- as.factor(national_case_final$Party)
national_case_final$Party <- factor(national_case_final$Party, 
                                    levels = rev(levels(national_case_final$Party)))
national_case_final <- national_case_final[which(national_case_final$Party != 'Other'), ]

ggplot(national_case_final, aes(x = X, 
                                y = Y_linear, 
                                color = Party)) +
  theme_calc() +
  theme(text = element_text(size = 12.5), 
        plot.title = element_text(hjust = 0.5),
        plot.margin = grid::unit(c(5, 5, 5, 5), "mm"),
        legend.title = element_blank(),
        legend.position ='bottom',
        legend.direction ='horizontal') +
  geom_line() +
  scale_color_manual(values = c('darkred', 'dodgerblue4')) +
  geom_ribbon(aes(ymin = L_linear, 
                  ymax = U_linear), 
              linetype = 2, 
              alpha = 0.0) +
  xlab('\nWeekly Count of National New Cases (per 10k)') + 
  ylab('Count of COVID-19 Tweets\n') +
  geom_rug(data = df, 
           aes(x = national_case_pop), 
           inherit.aes = FALSE, 
           color = 'gray60') +
  scale_shape_cleveland(overlap=FALSE)

ggsave(paste0(rep_path, 'Fig2c.png'),
       dpi = 600,
       width = 5,
       height = 3.5,
       units = 'in')



# national death -------------------------

national_death_values <- seq(round(min(df$national_death_pop), 1),
                             round(max(df$national_death_pop), 1),
                             0.1)
national_death_pop_effect <- Effect(focal.predictors = c('national_death_pop','party_new'), 
                                    mod = main_all_inter,
                                    xlevels = list(national_death_pop = national_death_values),
                                    given.values=c(
                                      state_abbrevCA = 0,
                                      state_abbrevTX = 0,
                                      state_abbrevAZ = 0,
                                      state_abbrevNY = 0,
                                      state_abbrevPA = 0,
                                      state_abbrevMA = 0,
                                      state_abbrevTN = 0,
                                      state_abbrevNM = 0,
                                      state_abbrevFL = 0,
                                      state_abbrevVA = 0,
                                      state_abbrevNC = 0,
                                      state_abbrevMN = 0,
                                      state_abbrevMD = 0,
                                      state_abbrevUT = 0,
                                      state_abbrevMI = 0,
                                      state_abbrevOH = 0,
                                      state_abbrevIA = 0,
                                      state_abbrevMO = 0,
                                      state_abbrevNJ = 0,
                                      state_abbrevNH = 0,
                                      state_abbrevWI = 0,
                                      state_abbrevAR = 0,
                                      state_abbrevOR = 0,
                                      state_abbrevGA = 0,
                                      state_abbrevCO = 1,
                                      state_abbrevSC = 0,
                                      state_abbrevCT = 0,
                                      state_abbrevOK = 0,
                                      state_abbrevAL = 0,
                                      state_abbrevDE = 0,
                                      state_abbrevNV = 0,
                                      state_abbrevKY = 0,
                                      state_abbrevIN = 0,
                                      state_abbrevND = 0,
                                      state_abbrevIL = 0,
                                      state_abbrevKS = 0,
                                      state_abbrevWA = 0,
                                      state_abbrevMS = 0,
                                      state_abbrevWV = 0,
                                      state_abbrevLA = 0,
                                      state_abbrevRI = 0,
                                      state_abbrevSD = 0,
                                      state_abbrevWY = 0,
                                      state_abbrevID = 0,
                                      state_abbrevHI = 0,
                                      state_abbrevMT = 0,
                                      state_abbrevME = 0,
                                      state_abbrevVT = 0,
                                      state_case_pop = median(df$state_case_pop),
                                      state_death_pop = median(df$state_death_pop),
                                      national_case_pop = median(df$national_case_pop),
                                      state_covid_policy = median(df$state_covid_policy),
                                      week_new = median(df$week_new),
                                      week_new_2 = median(df$week_new_2),
                                      week_new_3 = median(df$week_new_3)))

national_death_me_se_dem <- sqrt(vcov[7, 7]) # SE for marginal effect of predictor
national_death_me_se_other <- sqrt(vcov[7, 7] 
                                   + (1) ^ 2 * vcov[66, 66] + 
                                     2 * 1 * vcov[7, 66]) # SE for marginal effect of predictor
national_death_me_se_rep <- sqrt(vcov[7, 7] 
                                 + (1) ^ 2 * vcov[67, 67] + 
                                   2 * 1 * vcov[7, 67]) # SE for marginal effect of predictor

national_death_me_ci_dem <- cbind(national_death_values * (1.96 * national_death_me_se_dem), 
                                  national_death_values * -(1.96 * national_death_me_se_dem))
national_death_me_ci_other <- cbind(national_death_values * (1.96 * national_death_me_se_other), 
                                    national_death_values * -(1.96 * national_death_me_se_other))
national_death_me_ci_rep <- cbind(national_death_values * (1.96 * national_death_me_se_rep), 
                                  national_death_values * -(1.96 * national_death_me_se_rep))

national_death_preds <- national_death_pop_effect$fit # in the order of Dem, Neither, Rep
national_death_upper <- national_death_pop_effect$fit + c(national_death_me_ci_dem[, 1], 
                                                          national_death_me_ci_other[, 1], national_death_me_ci_rep[, 1])
national_death_lower <- national_death_pop_effect$fit + c(national_death_me_ci_dem[, 2], 
                                                          national_death_me_ci_other[, 2], national_death_me_ci_rep[, 2])
national_death_final <- data.frame(Party = c(rep('Democrat',
                                                 length(national_death_values)), 
                                             rep('Other',
                                                 length(national_death_values)),
                                             rep('Republican',
                                                 length(national_death_values))), 
                                   X = as.numeric(national_death_values), 
                                   Y = as.numeric(national_death_preds), 
                                   U = as.numeric(national_death_upper), 
                                   L = as.numeric(national_death_lower),
                                   Y_linear = exp(as.numeric(national_death_preds)) - 1, 
                                   U_linear = exp(as.numeric(national_death_upper)) - 1, 
                                   L_linear = exp(as.numeric(national_death_lower)) - 1)

national_death_final$Party <- as.factor(national_death_final$Party)
national_death_final$Party <- factor(national_death_final$Party, 
                                     levels = rev(levels(national_death_final$Party)))
national_death_final <- national_death_final[which(national_death_final$Party!='Other'), ]

ggplot(national_death_final, 
       aes(x = X, 
       y = Y_linear, 
       color = Party)) + 
  theme_calc() +
  theme(text = element_text(size = 12.5), 
        plot.title = element_text(hjust = 0.5),
        plot.margin = grid::unit(c(5, 5, 5, 5), "mm"),
        legend.title = element_blank(),
        legend.position='bottom',
        legend.direction='horizontal') +
  geom_line() +
  scale_color_manual(values = c('darkred', 'dodgerblue4')) +
  geom_ribbon(aes(ymin = 
                    L_linear, 
                  ymax = U_linear), 
              linetype = 2, alpha = 0.0) +
  xlab('\nWeekly Count of National New Deaths (per 10k)') + 
  ylab('Count of COVID-19 Tweets\n') +
  geom_rug(data=df, 
           aes(x = national_death_pop), 
           inherit.aes = FALSE, 
           color = 'gray60') +
  scale_shape_cleveland(overlap = FALSE)

ggsave(paste0(rep_path, 'Fig2d.png'),
       dpi = 600,
       width = 5,
       height = 3.5,
       units = 'in')


# state policy -------------------------

state_policy_values <- seq(round(min(df$state_covid_policy), 1),
                           round(max(df$state_covid_policy), 1),
                           0.01)
state_policy_effect <- Effect(focal.predictors = c('state_covid_policy', 'party_new'), 
                              mod = main_all_inter,
                              xlevels = list(state_covid_policy = state_policy_values),
                              given.values=c(
                                state_abbrevCA = 0,
                                state_abbrevTX = 0,
                                state_abbrevAZ = 0,
                                state_abbrevNY = 0,
                                state_abbrevPA = 0,
                                state_abbrevMA = 0,
                                state_abbrevTN = 0,
                                state_abbrevNM = 0,
                                state_abbrevFL = 0,
                                state_abbrevVA = 0,
                                state_abbrevNC = 0,
                                state_abbrevMN = 0,
                                state_abbrevMD = 0,
                                state_abbrevUT = 0,
                                state_abbrevMI = 0,
                                state_abbrevOH = 0,
                                state_abbrevIA = 0,
                                state_abbrevMO = 0,
                                state_abbrevNJ = 0,
                                state_abbrevNH = 0,
                                state_abbrevWI = 0,
                                state_abbrevAR = 0,
                                state_abbrevOR = 0,
                                state_abbrevGA = 0,
                                state_abbrevCO = 1,
                                state_abbrevSC = 0,
                                state_abbrevCT = 0,
                                state_abbrevOK = 0,
                                state_abbrevAL = 0,
                                state_abbrevDE = 0,
                                state_abbrevNV = 0,
                                state_abbrevKY = 0,
                                state_abbrevIN = 0,
                                state_abbrevND = 0,
                                state_abbrevIL = 0,
                                state_abbrevKS = 0,
                                state_abbrevWA = 0,
                                state_abbrevMS = 0,
                                state_abbrevWV = 0,
                                state_abbrevLA = 0,
                                state_abbrevRI = 0,
                                state_abbrevSD = 0,
                                state_abbrevWY = 0,
                                state_abbrevID = 0,
                                state_abbrevHI = 0,
                                state_abbrevMT = 0,
                                state_abbrevME = 0,
                                state_abbrevVT = 0,
                                state_case_pop = median(df$state_case_pop),
                                state_death_pop = median(df$state_death_pop),
                                national_case_pop = median(df$national_case_pop),
                                national_death_pop = median(df$national_death_pop),
                                week_new = median(df$week_new),
                                week_new_2 = median(df$week_new_2),
                                week_new_3 = median(df$week_new_3)))

state_policy_me_se_dem <- sqrt(vcov[8, 8]) # SE for marginal effect of predictor
state_policy_me_se_other <- sqrt(vcov[8, 8] 
                                 + (1) ^ 2 * vcov[68, 68] + 
                                   2 * 1 * vcov[8, 68]) # SE for marginal effect of predictor
state_policy_me_se_rep <- sqrt(vcov[8, 8] 
                               + (1) ^ 2 * vcov[69, 69] + 
                                 2 * 1 * vcov[8, 69]) # SE for marginal effect of predictor

state_policy_me_ci_dem <- cbind(state_policy_values * (1.96 * state_policy_me_se_dem), 
                                state_policy_values * -(1.96*state_policy_me_se_dem))
state_policy_me_ci_other <- cbind(state_policy_values * (1.96 * state_policy_me_se_other), 
                                  state_policy_values * -(1.96*state_policy_me_se_other))
state_policy_me_ci_rep <- cbind(state_policy_values * (1.96 * state_policy_me_se_rep), 
                                state_policy_values * -(1.96 * state_policy_me_se_rep))

state_policy_preds <- state_policy_effect$fit # in the order of Dem, Neither, Rep
state_policy_upper <- state_policy_effect$fit + c(state_policy_me_ci_dem[, 1], 
                                                  state_policy_me_ci_other[, 1], 
                                                  state_policy_me_ci_rep[, 1])
state_policy_lower <- state_policy_effect$fit + c(state_policy_me_ci_dem[, 2], 
                                                  state_policy_me_ci_other[, 2], 
                                                  state_policy_me_ci_rep[, 2])
state_policy_final <- data.frame(Party = c(rep('Democrat',
                                               length(state_policy_values)), 
                                           rep('Other',
                                               length(state_policy_values)),
                                           rep('Republican',
                                               length(state_policy_values))), 
                                 X = as.numeric(state_policy_values), 
                                 Y = as.numeric(state_policy_preds), 
                                 U = as.numeric(state_policy_upper), 
                                 L = as.numeric(state_policy_lower),
                                 Y_linear = exp(as.numeric(state_policy_preds)) - 1, 
                                 U_linear = exp(as.numeric(state_policy_upper)) - 1, 
                                 L_linear = exp(as.numeric(state_policy_lower)) - 1)

state_policy_final$Party <- as.factor(state_policy_final$Party)
state_policy_final$Party <- factor(state_policy_final$Party, 
                                   levels = rev(levels(state_policy_final$Party)))
state_policy_final <- state_policy_final[which(state_policy_final$Party!='Other'), ]

ggplot(state_policy_final, aes(x = X, 
                               y = Y_linear, 
                               color = Party)) + 
  theme_calc() +
  theme(text = element_text(size = 12.5), 
        plot.title = element_text(hjust = 0.5),
        plot.margin = grid::unit(c(5, 5, 5, 5), "mm"),
        legend.title = element_blank(),
        legend.position='bottom',
        legend.direction='horizontal') +
  scale_x_continuous(breaks = seq(0, 11, by = 1))+
  scale_y_continuous(limits = c(0, 1.5)) +
  geom_line() +
  scale_color_manual(values = c('darkred', 'dodgerblue4')) +
  geom_ribbon(aes(ymin = L_linear, 
                  ymax = U_linear), 
              linetype = 2, 
              alpha = 0.0) +
  xlab('\nWeekly Count of State New COVID-19 Policies') + 
  ylab('Count of COVID-19 Tweets\n') +
  geom_histogram(
    data = df, 
    aes(x = state_covid_policy, y = ..density..), 
    binwidth = 1,
    inherit.aes = FALSE, 
    fill = 'gray60', 
    alpha = 0.35) +
  scale_shape_cleveland(overlap = FALSE)

ggsave(paste0(rep_path, 'Fig2e.png'),
       dpi = 600,
       width = 5,
       height = 3.5,
       units = 'in')

